from qiskit import QuantumCircuit, transpile
from qiskit_aer import AerSimulator
from qiskit.circuit.library import MCXGate
import numpy as np

def grover_operator(
    oracle: QuantumCircuit,
    state_preparation: QuantumCircuit | None = None,
    reflection_qubits: list[int] | None = None,
    insert_barriers: bool = False,
    name: str = "Q",
) -> QuantumCircuit:
    """
    构造通用 Grover 算子 Q = A S0 A^dagger Sf

    参数:
        oracle: 相位 Oracle 电路（对好态翻转相位）
        state_preparation: 状态准备电路 A，如果为 None 默认全 H
        reflection_qubits: 零反射 S0 作用的 qubits，默认全部
        insert_barriers: 是否在各阶段插入 barrier
        name: 电路名称

    返回:
        QuantumCircuit: Grover 算子电路
    """
    # 创建一个空电路，继承 oracle 的量子寄存器
    circuit = oracle.copy_empty_like(name=name)

    # (1) 加入 oracle
    circuit.compose(oracle, inplace=True)
    if insert_barriers:
        circuit.barrier()

    # (2) 逆的状态准备 A^dagger
    if reflection_qubits is None:
        reflection_qubits = list(range(circuit.num_qubits))

    if state_preparation is None:
        circuit.h(reflection_qubits)
    else:
        circuit.compose(state_preparation.inverse(), inplace=True)

    if insert_barriers:
        circuit.barrier()

    # (3) 零反射 S0
    n_reflect = len(reflection_qubits)
    circuit.x(reflection_qubits)  # 先 X 翻转，把 |0> 映射到 |1>
    if n_reflect == 1:
        circuit.z(reflection_qubits[0])
    else:
        mcx = MCXGate(n_reflect - 1)  # 多控 X，控制位 = n-1
        circuit.h(reflection_qubits[-1])
        circuit.append(mcx, reflection_qubits)
        circuit.h(reflection_qubits[-1])
    circuit.x(reflection_qubits)  # 恢复原态

    if insert_barriers:
        circuit.barrier()

    # (4) 再次状态准备 A
    if state_preparation is None:
        circuit.h(reflection_qubits)
    else:
        circuit.compose(state_preparation, inplace=True)

    # (5) 全局相位 -1
    circuit.global_phase = np.pi

    return circuit

# =====================
# 测试 Grover 算子有效性
# =====================
if __name__ == "__main__":
    # 2 qubit，目标态 |01>
    oracle = QuantumCircuit(2)
    # 构造多比特 Oracle，只翻转 |01> 的相位
    # X 翻转第一个 qubit -> |01> 变成 |11>
    oracle.x(0)
    # MCZ (这里用 CZ) 翻转 |11>
    oracle.cz(0, 1)
    # 恢复第一个 qubit
    oracle.x(0)
    oracle.barrier()
    # 构造 Grover 算子
    grover_op = grover_operator(oracle)

    # 构造完整搜索电路
    qc = QuantumCircuit(2, 2)
    qc.h([0, 1])  # 初始化均匀叠加态
    qc.compose(grover_op, [0, 1], inplace=True)
    qc.measure([0, 1], [0, 1])
    print(qc)
    # 仿真
    simulator = AerSimulator()
    result = simulator.run(qc, shots=1024).result()
    counts = result.get_counts()

    print("测量结果统计:")
    print(counts)
